package demo.huodongbaoming.loadbalancer;

import org.springframework.beans.factory.ObjectProvider;
import org.springframework.cloud.client.ServiceInstance;
import org.springframework.cloud.client.loadbalancer.DefaultResponse;
import org.springframework.cloud.client.loadbalancer.EmptyResponse;
import org.springframework.cloud.client.loadbalancer.Request;
import org.springframework.cloud.client.loadbalancer.Response;
import org.springframework.cloud.loadbalancer.core.ReactorServiceInstanceLoadBalancer;
import org.springframework.cloud.loadbalancer.core.ServiceInstanceListSupplier;
import reactor.core.publisher.Mono;

import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * 自定义负载均衡器
 * 实现加权轮询策略，可以根据服务实例的元数据中的权重进行负载均衡
 */
public class CustomLoadBalancer implements ReactorServiceInstanceLoadBalancer {

    private final String serviceId;
    private final ObjectProvider<ServiceInstanceListSupplier> serviceInstanceListSupplierProvider;
    private final Map<String, AtomicInteger> positionMap = new ConcurrentHashMap<>();

    public CustomLoadBalancer(ObjectProvider<ServiceInstanceListSupplier> serviceInstanceListSupplierProvider, 
                             String serviceId) {
        this.serviceId = serviceId;
        this.serviceInstanceListSupplierProvider = serviceInstanceListSupplierProvider;
    }

    @Override
    public Mono<Response<ServiceInstance>> choose(Request request) {
        ServiceInstanceListSupplier supplier = serviceInstanceListSupplierProvider.getIfAvailable();
        return supplier.get().next().map(this::processInstanceResponse);
    }

    private Response<ServiceInstance> processInstanceResponse(List<ServiceInstance> instances) {
        if (instances.isEmpty()) {
            return new EmptyResponse();
        }
        
        // 获取当前服务的计数器，如果不存在则创建
        AtomicInteger position = positionMap.computeIfAbsent(serviceId, k -> new AtomicInteger(0));
        
        // 如果只有一个实例，直接返回
        if (instances.size() == 1) {
            return new DefaultResponse(instances.get(0));
        }

        // 计算下一个位置，并确保不超出列表范围
        int pos = position.incrementAndGet() % instances.size();
        if (pos < 0) {
            pos = 0;
            position.set(0);
        }
        
        // 返回选择的服务实例
        return new DefaultResponse(instances.get(pos));
    }
} 